/*
Copyright (c) 2025 Huawei Technologies Co., Ltd.
This file is a part of the CANN Open Software.
Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
Please refer to the License for details. You may not use this file except in compliance with the License.
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
See LICENSE in the root of the software repository for the full text of the License. 
*/

#pragma once
#include "acl/acl.h"
#include "../include/custom_type.h"
#include "./data_utils.h"
#include "../include/kernel_const.h"

#include <algorithm>
#include <cstdlib>
#include <cstdio>
#include <time.h>
#include <sys/time.h>  // 用于Linux系统


#ifdef __cplusplus
extern "C" aclError rtGetC2cCtrlAddr(uint64_t* addr, uint32_t* len);
#endif


uint8_t printTest = 0; 

/**
 * @brief: 通过掩码矩阵计算有效长度
 * @param[in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param[in] batchCount: 批量矩阵乘batch数
 * @param[in] h_maskA: host侧掩码矩阵空间
 * @param[out] d_maskA: device侧掩码矩阵空间
 * @param[out] h_validM: 通过掩码矩阵计算得有效长度
 * @param[in] mode: 算子任务选择
*/
void calHValidMByMaskA(
    uint32_t zeroPaddingM, 
    uint32_t batchCount, 
    half *h_maskA, 
    half *d_maskA, 
    uint32_t *h_validM, 
    uint8_t mode
){
    
    if(mode == 0){

        uint64_t inputMaskABytes = sizeof(half) * batchCount * zeroPaddingM; 
        std::fill(h_validM, h_validM + batchCount, zeroPaddingM); 
        std::fill(h_maskA, h_maskA + batchCount * zeroPaddingM, 1.0f); 
        ACL_CHECK(aclrtMemcpy(d_maskA, inputMaskABytes, h_maskA, inputMaskABytes, ACL_MEMCPY_HOST_TO_DEVICE)); 

    }else{

        for(uint32_t i = 0; i < batchCount; i++){
            half *curMask = h_maskA + i * zeroPaddingM;
            uint32_t left = 0, right = zeroPaddingM; 
            uint32_t mid; 
            while(right - left > 1){
                mid = left + (right - left) / 2;
                if(curMask[mid] != 0.0f){
                    left = mid;
                }else{
                    right = mid;
                }
            }
            h_validM[i] = left + 1;
        }
    }

    if(printTest){
        printf("Host cal validM: \n");
        for(uint64_t i = 0; i < batchCount; i++ ){
            printf("%d ", h_validM[i]);
        }
        printf("\n");
    }
    
    
}

/**
 * @brief: 无掩码矩阵计算有效长度
 * @param[in] layoutA: A矩阵排布格式
 * @param[in] h_A: 连续存储的零填充A矩阵首地址数组
 * @param[in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param[in] zeroPaddingK: A、B矩阵零填充后的K维度
 * @param[in] batchCount: 批量矩阵乘batch数
 * @param[out] h_validM: 通过掩码矩阵计算得有效长度
*/
void calHValidM(
    layoutType layoutA, 
    half* h_A, 
    uint32_t zeroPaddingM, 
    uint32_t zeroPaddingK, 
    uint32_t batchCount, 
    uint32_t* h_validM
){

    struct timeval start, end;
    gettimeofday(&start, NULL);

    uint32_t strideA = zeroPaddingM * zeroPaddingK;
    for(uint32_t i = 0; i < batchCount; i++){
        half* h_A_cur = h_A + i * strideA;

        uint32_t left = 0, right = zeroPaddingM;
        uint32_t mid;
        while(right - left > 1){
            mid = left + (right - left) / 2;
            if( h_A_cur[mid * (layoutA == RowMajor ? zeroPaddingK : 1)] != 0.0f ){
                left = mid;
            }else{
                right = mid;
            }
        }
        h_validM[i] = left + 1;

    }

    gettimeofday(&end, NULL);
    uint32_t elapsed = (end.tv_sec - start.tv_sec) * 1000000L + (end.tv_usec - start.tv_usec);
    printf("Elapsed time of calc: %d us. \n", elapsed);

    // test
    printf("The calculated validM: \n");
    for(uint32_t i = 0; i < batchCount; i++ ){
        printf("%d ", h_validM[i]);
    }
    printf("\n\n");
}


/**
 * @brief: 通过任务选择填充其余维度有效长度
 * @param[in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param[in] zeroPaddingN: B、C矩阵零填充后的N维度
 * @param[in] zeroPaddingK: A、B矩阵零填充后的K维度
 * @param[in] batchCount: 批量矩阵乘batch数
 * @param[out] h_validN: 通过任务选择填充N维度有效长度
 * @param[out] h_validM: 通过任务选择填充K维度有效长度
 * @param[in] mode: 算子任务选择
*/
void calHValidNK(
    uint32_t zeroPaddingM, 
    uint32_t zeroPaddingN, 
    uint32_t zeroPaddingK, 
    uint32_t batchCount, 
    uint32_t* h_validM, 
    uint32_t* h_validN, 
    uint32_t* h_validK, 
    uint8_t mode
){

    uint64_t inputValidMBytes = sizeof( uint32_t ) * batchCount;
    uint64_t inputValidNBytes = sizeof( uint32_t ) * batchCount; 
    uint64_t inputValidKBytes = sizeof( uint32_t ) * batchCount;

    if(mode == 0){
        std::fill(h_validN, h_validN + batchCount, zeroPaddingN);
        std::fill(h_validK, h_validK + batchCount, zeroPaddingK);
    }else if(mode == 1){
        ACL_CHECK(aclrtMemcpy(h_validN, inputValidNBytes, h_validM, inputValidMBytes, ACL_MEMCPY_HOST_TO_HOST));
        std::fill(h_validK, h_validK + batchCount, zeroPaddingK);
    }else if(mode == 2){
        std::fill(h_validN, h_validN + batchCount, zeroPaddingN);
        ACL_CHECK(aclrtMemcpy(h_validK, inputValidKBytes, h_validM, inputValidMBytes, ACL_MEMCPY_HOST_TO_HOST));
    }else if(mode == 3){
        std::fill(h_validN, h_validN + batchCount, zeroPaddingN);
        std::fill(h_validK, h_validK + batchCount, zeroPaddingK);
    }   

    if(printTest){
        printf("Host cal validN: \n");
        for(uint32_t i = 0; i < batchCount; i++ ){
            printf("%d ", h_validN[i]);
        }
        printf("\n");
        printf("Host cal validK: \n");
        for(uint32_t i = 0; i < batchCount; i++ ){
            printf("%d ", h_validK[i]);
        }
        printf("\n");
    }
    
}


/**
 * @brief: 计算不同batch矩阵首地址
 * @param[in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param[in] zeroPaddingN: B、C矩阵零填充后的N维度
 * @param[in] zeroPaddingK: A、B矩阵零填充后的K维度
 * @param[in] h_A: 连续存储的零填充A矩阵首地址数组
 * @param[in] h_B: 连续存储的零填充B矩阵首地址数组
 * @param[in] h_B: 连续存储的零填充B矩阵首地址数组
 * @param[in] batchCount: 批量矩阵乘batch数
 * @param[out] h_APointer: 不同batch A矩阵首地址
 * @param[out] h_BPointer: 不同batch B矩阵首地址
 * @param[out] h_CPointer: 不同batch C矩阵首地址
*/
void calHPointer(
    uint32_t zeroPaddingM, 
    uint32_t zeroPaddingN, 
    uint32_t zeroPaddingK, 
    half *d_A,  
    half *d_B, 
    half *d_C, 
    uint32_t batchCount, 
    half** h_APointer, 
    half** h_BPointer, 
    half** h_CPointer
){
    uint32_t strideSingleA = zeroPaddingM * zeroPaddingK;
    uint32_t strideSingleB = zeroPaddingK * zeroPaddingN;
    uint32_t strideSingleC = zeroPaddingM * zeroPaddingN;
    h_APointer[0] = d_A;
    h_BPointer[0] = d_B;
    h_CPointer[0] = d_C;
    for(uint32_t i = 1; i < batchCount; i++ ) {
         h_APointer[i] = h_APointer[i-1] + strideSingleA;
         h_BPointer[i] = h_BPointer[i-1] + strideSingleB;
         h_CPointer[i] = h_CPointer[i-1] + strideSingleC;
    }
    
}

/**
 * @brief: 对ABC矩阵进行任务分块的分块大小进行赋值
 * @param[out] L1_M0
 * @param[out] L1_N0
 * @param[out] L1_K0
*/
void calL1M0N0K0(
    uint32_t &L1_M0, 
    uint32_t &L1_N0, 
    uint32_t &L1_K0
){
    L1_M0 = L1M0; 
    L1_N0 = L1N0; 
    L1_K0 = L1K0; 

    if(printTest){
        printf("L1_M0 = %d, L1_N0 = %d, L1_K0 = %d. \n", L1_M0, L1_N0, L1_K0);
    }
}

/**
 * @brief: 计算不同A、B矩阵是否需要padding
 * @param[in] layoutA: A矩阵排布格式
 * @param[in] layoutB: B矩阵排布格式
 * @param[in] batchCount: 批量矩阵乘batch数
 * @param[in] L1_M0: A、C矩阵任务分块的M维度大小
 * @param[in] L1_N0: B、C矩阵任务分块的N维度大小
 * @param[in] L1_K0: A、B矩阵任务分块的K维度大小
 * @param[in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param[in] zeroPaddingN: B、C矩阵零填充后的N维度
 * @param[in] zeroPaddingK: A、B矩阵零填充后的K维度
 * @param[out] h_isAPadding: 各个batch A矩阵是否需要padding
 * @param[out] h_isBPadding: 各个batch B矩阵是否需要padding
*/
void calIsPaddingAndLda(
    layoutType layoutA, 
    layoutType layoutB, 
    uint32_t batchCount, 
    uint32_t L1_M0, 
    uint32_t L1_N0, 
    uint32_t L1_K0, 
    uint32_t zeroPaddingM, 
    uint32_t zeroPaddingN, 
    uint32_t zeroPaddingK, 
    uint8_t *h_isAPadding, 
    uint8_t *h_isBPadding
){
    for(uint32_t i = 0; i < batchCount; i++ ){
        uint32_t AMajorSize = ( layoutA == RowMajor ? zeroPaddingK : zeroPaddingM );
        uint32_t BMajorSize = ( layoutB == RowMajor ? zeroPaddingN : zeroPaddingK );
        uint32_t AMajorPaddingSize = ( layoutA == RowMajor ? L1_K0 : L1_M0 );
        uint32_t BMajorPaddingSize = ( layoutB == RowMajor ? L1_N0 : L1_K0 );
        if( AMajorSize % AMajorPaddingSize != 0){
            h_isAPadding[i] = 1;
        }else{
            h_isAPadding[i] = 0;
        }
        if( BMajorSize % BMajorPaddingSize != 0){
            h_isBPadding[i] = 1;
        }else{
            h_isBPadding[i] = 0;
        }
    }

    if(printTest){
        printf("h_isAPadding: \n");
        for(uint64_t i = 0; i < batchCount; i++ ){
            printf("%s ", h_isAPadding[i]==1? "Y" : "N");
        }
        printf("\n");

        printf("h_isBPadding: \n");
        for(uint64_t i = 0; i < batchCount; i++ ){
            printf("%s ", h_isBPadding[i]==1? "Y" : "N");
        }
        printf("\n");
    }
    
}

void genPaddingDir(
    layoutType layoutA, 
    layoutType layoutB, 
    uint8_t &paddingDirA, 
    uint8_t &paddingDirB
){
    if(layoutA == RowMajor){
        paddingDirA = 1; 
    }else{
        paddingDirA = 2; 
    }
    if(layoutB == RowMajor){
        paddingDirB = 2; 
    }else{
        paddingDirB = 1; 
    }
    if(printTest){
        printf("PaddingDirA = %d, PaddingDirB = %d. \n", paddingDirA, paddingDirB); 
    }
    
}

/**
 * @brief: 为需要padding的矩阵开辟新的device内存空间
 * @param[in] batchCount: 批量矩阵乘batch数
 * @param[in] L1_M0: A、C矩阵任务分块的M维度大小
 * @param[in] L1_N0: B、C矩阵任务分块的N维度大小
 * @param[in] L1_K0: A、B矩阵任务分块的K维度大小
 * @param[in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param[in] zeroPaddingN: B、C矩阵零填充后的N维度
 * @param[in] zeroPaddingK: A、B矩阵零填充后的K维度
 * @param[in] h_isAPadding: 各个batch A矩阵是否需要padding
 * @param[in] h_isBPadding: 各个batch B矩阵是否需要padding
 * @param[out] h_APointerPadding: 需padding的A矩阵新的device内存空间
 * @param[out] h_BPointerPadding: 需padding的B矩阵新的device内存空间
*/
void allocPointerPadding(
    uint32_t batchCount, 
    uint32_t L1_M0, 
    uint32_t L1_N0, 
    uint32_t L1_K0, 
    uint32_t zeroPaddingM, 
    uint32_t zeroPaddingN, 
    uint32_t zeroPaddingK, 
    uint8_t *h_isAPadding, 
    uint8_t *h_isBPadding, 
    half **h_APointerPadding,  
    half **h_BPointerPadding
){
    for(uint64_t i = 0; i < batchCount; i++ ){
        if( h_isAPadding[i] ){
            half *d_A_pointer_padding_i = nullptr;
            uint64_t stridePaddingABytes = sizeof(half) * RoundUp<uint32_t>(zeroPaddingM, L1_M0) * RoundUp<uint32_t>(zeroPaddingK, L1_K0);
            ACL_CHECK(aclrtMalloc((void **)(&d_A_pointer_padding_i), stridePaddingABytes, ACL_MEM_MALLOC_HUGE_FIRST));
            h_APointerPadding[i] = d_A_pointer_padding_i;
        }else{
            h_APointerPadding[i] = nullptr;
        }
        if( h_isBPadding[i] ){
            half *d_B_pointer_padding_i = nullptr;
            uint64_t stridePaddingBBytes = sizeof(half) * RoundUp<uint32_t>(zeroPaddingN, L1_N0) * RoundUp<uint32_t>(zeroPaddingK, L1_K0);
            ACL_CHECK(aclrtMalloc((void **)(&d_B_pointer_padding_i), stridePaddingBBytes, ACL_MEM_MALLOC_HUGE_FIRST));
            h_BPointerPadding[i] = d_B_pointer_padding_i;
        }else{
            h_BPointerPadding[i] = nullptr;
        }
    }

    if(printTest){
        printf("h_APointerPadding: \n");
        for(uint64_t i = 0; i < batchCount; i++){
            printf("%s ", h_APointerPadding[i] == nullptr? "Null" : "NotNull");
        }
        printf("\n");

        printf("h_BPointerPadding: \n");
        for(uint64_t i = 0; i < batchCount; i++){
            printf("%s ", h_BPointerPadding[i] == nullptr? "Null" : "NotNull");
        }
        printf("\n");
    }

}

/**
 * @brief: 计算C矩阵任务分块数
 * @param[in] batchCount: 批量矩阵乘batch数
 * @param[out] h_CBlocks: C矩阵任务分块数
 * @param[out] cBlocksSum: C矩阵任务分块总数
 * @param[in] h_validN: 通过任务选择填充N维度有效长度
 * @param[in] h_validM: 通过任务选择填充K维度有效长度
 * @param[in] L1_M0: A、C矩阵任务分块的M维度大小
 * @param[in] L1_N0: B、C矩阵任务分块的N维度大小
*/
void calcBlocks(
    uint32_t batchCount, 
    uint32_t *h_CBlocks, 
    uint32_t &cBlocksSum, 
    uint32_t *h_validM, 
    uint32_t *h_validN, 
    uint32_t L1_M0, 
    uint32_t L1_N0
){

    for(uint64_t i = 0; i < batchCount; i++ ){
        h_CBlocks[i] = CeilDiv<uint32_t>(h_validM[i], L1_M0) * CeilDiv<uint32_t>(h_validN[i], L1_N0);
        cBlocksSum += h_CBlocks[i];
    }

    if(printTest){
        printf("h_CBlocks: \n");
        for(uint64_t i = 0; i < batchCount; i++){
            printf("%d ", h_CBlocks[i]);
        }
        printf("\ncBlocksSum = %d\n", cBlocksSum);
    }
    
}

/**
 * @brief: 计算并分配workspace
 * @param[out] h_AicAivWorkspacePointer
 * @param[out] d_AicAivWorkspace
 * @param[in] L1_M0: A、C矩阵任务分块的M维度大小
 * @param[in] L1_N0: B、C矩阵任务分块的N维度大小
*/
void allocWorkspace(
    half** &h_AicAivWorkspacePointer,
    half* &d_AicAivWorkspace, 
    uint32_t L1_M0, 
    uint32_t L1_N0
){

    uint64_t WORKSPACE_NUM = WORKSPACENUM;
    uint64_t workspaceSingleCoreSize = WORKSPACE_NUM * L1_M0 * L1_N0;
    uint64_t workspaceBytes = MAX_AICORE_NUM * workspaceSingleCoreSize * sizeof(half);
    ACL_CHECK(aclrtMalloc((void**)(&d_AicAivWorkspace), workspaceBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    h_AicAivWorkspacePointer[0] = d_AicAivWorkspace;
    for(uint64_t i = 1; i < /*batchCount*/ MAX_AICORE_NUM; i++){
        h_AicAivWorkspacePointer[i] = h_AicAivWorkspacePointer[i-1] + workspaceSingleCoreSize;
    }
}

/**
 * @brief: 打印输出kernel计算的有效长度
 * @param[in] batchCount: 批量矩阵乘batch数
 * @param[in] h_validM: M维度有效长度
 * @param[in] h_validN: N维度有效长度
 * @param[in] h_validK: K维度有效长度
*/
void printCalValidLen(
    uint32_t batchCount, 
    uint32_t *h_validM, 
    uint32_t *h_validN, 
    uint32_t *h_validK
){
    if(printTest){
        printf("Kernel cal validM: \n"); 
        for(uint32_t i = 0; i < batchCount; i++){
            printf("%d ", h_validM[i]);
        }
        printf("\n");
        printf("Kernel cal validN: \n"); 
        for(uint32_t i = 0; i < batchCount; i++){
            printf("%d ", h_validN[i]);
        }
        printf("\n");
        printf("Kernel cal validK: \n"); 
        for(uint32_t i = 0; i < batchCount; i++){
            printf("%d ", h_validK[i]);
        }
        printf("\n");
    }
    
}

/**
 * @brief: 
 * @param [in] kernel_device: 要调用的算子
 * @param [in] layoutA: A矩阵排布格式
 * @param [in] layoutB: B矩阵排布格式
 * @param [in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param [in] zeroPaddingN: B、C矩阵零填充后的N维度
 * @param [in] zeroPaddingK: A、B矩阵零填充后的K维度
 * @param [in] batchCount: 批量矩阵乘的batch数
 * @param [in] d_maskA: 掩码矩阵
 * @param [in] alpha: alpha*AB+beta*C
 * @param [in] d_A: 连续存储的零填充A矩阵首地址数组
 * @param [in] d_B: 连续存储的零填充B矩阵首地址数组
 * @param [in] beta:  alpha*AB+beta*C
 * @param [out] d_C: 每批矩阵乘的零填充B矩阵首地址数组
 * @param [in] stream: ACL同步流
 * @param [in] mode: 算子任务选择
 */
void kernel_host (
    void (*kernel_device)(
        layoutType,             // layoutA
        layoutType,             // layoutB
        uint32_t,               // zeroPaddingM
        uint32_t,               // zeroPaddingN
        uint32_t,               // zeroPaddingK
        uint32_t,               // batchCount
        __gm__ half *,          // d_maskA
        __gm__ uint32_t*,       // d_validM
        __gm__ uint32_t*,       // d_validN
        __gm__ uint32_t*,       // d_validK
        half,                   // alpha
        __gm__ half**,          // d_APointer
        __gm__ half**,          // d_BPointer
        half,                   // beta
        __gm__ half**,          // d_CPointer
        __gm__ uint8_t*,        // d_isAPadding
        __gm__ uint8_t*,        // d_isBPadding
        __gm__ half**,          // d_APointerPadding
        __gm__ half**,          // d_BPointerPadding
        uint8_t,                // paddingDirA
        uint8_t,                // paddingDirB
        __gm__ half**,      // d_AicAivWorkspacePointer
        uint64_t,           // fftsAddr
        uint8_t             // is_alpha1_beta0
    ), 
    layoutType layoutA, 
    layoutType layoutB, 
    uint32_t zeroPaddingM, 
    uint32_t zeroPaddingN, 
    uint32_t zeroPaddingK, 
    uint32_t batchCount, 
    half *d_maskA, 
    half alpha,  
    half *d_A,  
    half *d_B, 
    half beta, 
    half *d_C,
    aclrtStream &stream, 
    uint8_t mode
) {

    INFO_LOG("Kernel host starts execution. ");

    uint64_t inputValidMBytes = sizeof( uint32_t ) * batchCount;
    uint64_t inputValidNBytes = sizeof( uint32_t ) * batchCount; 
    uint64_t inputValidKBytes = sizeof( uint32_t ) * batchCount;
    uint64_t inputMaskABytes = sizeof( half ) * zeroPaddingM * batchCount; 

    uint32_t *h_validM = nullptr;
    uint32_t *h_validN = nullptr;
    uint32_t *h_validK = nullptr;
    half *h_maskA = nullptr; 
    
    ACL_CHECK(aclrtMallocHost((void **)(&h_maskA), inputMaskABytes));
    ACL_CHECK(aclrtMemcpy(h_maskA, inputMaskABytes, d_maskA, inputMaskABytes, ACL_MEMCPY_DEVICE_TO_HOST));

    ACL_CHECK(aclrtMallocHost((void **)(&h_validM), inputValidMBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_validN), inputValidNBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_validK), inputValidKBytes));

    calHValidMByMaskA(
        zeroPaddingM, 
        batchCount, 
        h_maskA, 
        d_maskA, 
        h_validM, 
        mode
    );

    calHValidNK(
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        batchCount, 
        h_validM, 
        h_validN, 
        h_validK, 
        mode
    );

    uint32_t *d_validM = nullptr; 
    uint32_t *d_validN = nullptr;
    uint32_t *d_validK = nullptr;

    ACL_CHECK(aclrtMalloc((void **)(&d_validM), inputValidMBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_validN), inputValidNBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_validK), inputValidKBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    ACL_CHECK(aclrtMemset(d_validM, inputValidMBytes, 0, inputValidMBytes)); 

    half **h_APointer = nullptr;
    half **h_BPointer = nullptr;
    half **h_CPointer = nullptr;

    uint64_t ABCPointerBytes = batchCount * sizeof( half * );

    ACL_CHECK(aclrtMallocHost((void **)(&h_APointer), ABCPointerBytes)); 
    ACL_CHECK(aclrtMallocHost((void **)(&h_BPointer), ABCPointerBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_CPointer), ABCPointerBytes));

    calHPointer(
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        d_A, 
        d_B, 
        d_C, 
        batchCount, 
        h_APointer, 
        h_BPointer, 
        h_CPointer
    );

    half **d_APointer = nullptr;
    half **d_BPointer = nullptr;
    half **d_CPointer = nullptr;

    ACL_CHECK(aclrtMalloc((void **)(&d_APointer), ABCPointerBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_BPointer), ABCPointerBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_CPointer), ABCPointerBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    ACL_CHECK(aclrtMemcpy(d_APointer, ABCPointerBytes, h_APointer, ABCPointerBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_BPointer, ABCPointerBytes, h_BPointer, ABCPointerBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_CPointer, ABCPointerBytes, h_CPointer, ABCPointerBytes, ACL_MEMCPY_HOST_TO_DEVICE));

    uint32_t L1_M0 = L1M0;
    uint32_t L1_N0 = L1N0;
    uint32_t L1_K0 = L1K0;

    calL1M0N0K0(
        L1_M0, 
        L1_N0, 
        L1_K0
    );

    uint8_t *h_isAPadding = nullptr;
    uint8_t *h_isBPadding = nullptr;
    half **h_APointerPadding = nullptr;
    half **h_BPointerPadding = nullptr;
    uint64_t isPaddingBytes = batchCount * sizeof(uint8_t);
    uint64_t ABPointerPaddingBytes = batchCount * sizeof( half * );
    uint64_t validPaddingLdabBytes = batchCount * sizeof(uint32_t);

    ACL_CHECK(aclrtMallocHost((void **)(&h_isAPadding), isPaddingBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_isBPadding), isPaddingBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_APointerPadding), ABPointerPaddingBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_BPointerPadding), ABPointerPaddingBytes));

    calIsPaddingAndLda(
        layoutA, 
        layoutB, 
        batchCount, 
        L1_M0, 
        L1_N0, 
        L1_K0, 
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        h_isAPadding, 
        h_isBPadding
    );

    uint8_t *d_isAPadding = nullptr;
    uint8_t *d_isBPadding = nullptr;

    ACL_CHECK(aclrtMalloc((void **)(&d_isAPadding), isPaddingBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_isBPadding), isPaddingBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    ACL_CHECK(aclrtMemcpy(d_isAPadding, isPaddingBytes, h_isAPadding, isPaddingBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_isBPadding, isPaddingBytes, h_isBPadding, isPaddingBytes, ACL_MEMCPY_HOST_TO_DEVICE));

    allocPointerPadding(
        batchCount, 
        L1_M0, 
        L1_N0, 
        L1_K0, 
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        h_isAPadding, 
        h_isBPadding, 
        h_APointerPadding, 
        h_BPointerPadding
    );

    half **d_APointerPadding = nullptr;
    half **d_BPointerPadding = nullptr;

    ACL_CHECK(aclrtMalloc((void **)(&d_APointerPadding), ABPointerPaddingBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_BPointerPadding), ABPointerPaddingBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    ACL_CHECK(aclrtMemcpy(d_APointerPadding, ABPointerPaddingBytes, h_APointerPadding, ABPointerPaddingBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_BPointerPadding, ABPointerPaddingBytes, h_BPointerPadding, ABPointerPaddingBytes, ACL_MEMCPY_HOST_TO_DEVICE));

    uint32_t cBlocksSum = 0;
    uint32_t *h_CBlocks = nullptr;
    uint64_t cBlocksBytes = batchCount * sizeof(uint32_t);

    ACL_CHECK(aclrtMallocHost((void **)(&h_CBlocks), cBlocksBytes));
    calcBlocks(
        batchCount, 
        h_CBlocks, 
        cBlocksSum, 
        h_validM, 
        h_validN, 
        L1_M0, 
        L1_N0
    );

    
    half** h_AicAivWorkspacePointer = nullptr; 
    half* d_AicAivWorkspace = nullptr;
    uint64_t workspacePointerBytes = MAX_AICORE_NUM * sizeof(half *); 
    ACL_CHECK(aclrtMallocHost((void **)(&h_AicAivWorkspacePointer), workspacePointerBytes));

    allocWorkspace(
        h_AicAivWorkspacePointer, 
        d_AicAivWorkspace, 
        L1_M0, 
        L1_N0
    );

    half** d_AicAivWorkspacePointer = nullptr;
    ACL_CHECK(aclrtMalloc((void **)(&d_AicAivWorkspacePointer), workspacePointerBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMemcpy(d_AicAivWorkspacePointer, workspacePointerBytes, h_AicAivWorkspacePointer, workspacePointerBytes, ACL_MEMCPY_HOST_TO_DEVICE));


    uint64_t fftsAddr{0};
    uint32_t fftsLen{0};
    ACL_CHECK(rtGetC2cCtrlAddr(&fftsAddr, &fftsLen));

    uint8_t paddingDirA = 0; 
    uint8_t paddingDirB = 0; 
    genPaddingDir(
        layoutA, 
        layoutB, 
        paddingDirA, 
        paddingDirB
    );

    uint8_t is_alpha1_beta0 = (alpha == 1.0f && beta == 0.0f );

    struct timeval startKernel, endKernel;
    gettimeofday(&startKernel, NULL);
    INFO_LOG("Kernel device starts execution. ");
    (*kernel_device)<<< MAX_AICORE_NUM, nullptr, stream >>>(
        layoutA, 
        layoutB, 
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        batchCount, 
        d_maskA, 
        d_validM, 
        d_validN, 
        d_validK, 
        alpha, 
        d_APointer, 
        d_BPointer, 
        beta, 
        d_CPointer, 
        d_isAPadding, 
        d_isBPadding, 
        d_APointerPadding, 
        d_BPointerPadding, 
        paddingDirA, 
        paddingDirB, 
        d_AicAivWorkspacePointer, 
        fftsAddr, 
        is_alpha1_beta0
    );

    ACL_CHECK(aclrtSynchronizeStream(stream));
    gettimeofday(&endKernel, NULL);
    uint32_t elapsedKernel = (endKernel.tv_sec - startKernel.tv_sec) * 1000000L + (endKernel.tv_usec - startKernel.tv_usec);
    printf("Elapsed time of kernel: %d us. \n", elapsedKernel);
    INFO_LOG("Kernel device was executed successfully. ");

    ACL_CHECK(aclrtMemcpy(h_validM, inputValidMBytes, d_validM, inputValidMBytes, ACL_MEMCPY_DEVICE_TO_HOST));
    ACL_CHECK(aclrtMemcpy(h_validN, inputValidNBytes, d_validN, inputValidNBytes, ACL_MEMCPY_DEVICE_TO_HOST));
    ACL_CHECK(aclrtMemcpy(h_validK, inputValidKBytes, d_validK, inputValidKBytes, ACL_MEMCPY_DEVICE_TO_HOST));

    printCalValidLen(
        batchCount, 
        h_validM, 
        h_validN, 
        h_validK
    ); 

    ACL_CHECK(aclrtFreeHost(h_validM));
    ACL_CHECK(aclrtFreeHost(h_validN));
    ACL_CHECK(aclrtFreeHost(h_validK));
    ACL_CHECK(aclrtFreeHost(h_APointer));
    ACL_CHECK(aclrtFreeHost(h_BPointer));
    ACL_CHECK(aclrtFreeHost(h_CPointer));

    for(uint32_t i = 0; i < batchCount; i++){
        if( h_isAPadding[i] ){
            ACL_CHECK(aclrtFree(h_APointerPadding[i]));
        }
        if( h_isBPadding[i] ){
            ACL_CHECK(aclrtFree(h_BPointerPadding[i]));
        }
    }
    ACL_CHECK(aclrtFreeHost(h_isAPadding));
    ACL_CHECK(aclrtFreeHost(h_isBPadding));
    ACL_CHECK(aclrtFreeHost(h_APointerPadding));
    ACL_CHECK(aclrtFreeHost(h_BPointerPadding));

    ACL_CHECK(aclrtFreeHost(h_CBlocks));
    ACL_CHECK(aclrtFreeHost(h_AicAivWorkspacePointer));

    ACL_CHECK(aclrtFree(d_validM)); 
    ACL_CHECK(aclrtFree(d_validN));
    ACL_CHECK(aclrtFree(d_validK));
    ACL_CHECK(aclrtFree(d_APointer));
    ACL_CHECK(aclrtFree(d_BPointer));
    ACL_CHECK(aclrtFree(d_CPointer));

    ACL_CHECK(aclrtFree(d_isAPadding));
    ACL_CHECK(aclrtFree(d_isBPadding));
    ACL_CHECK(aclrtFree(d_APointerPadding));
    ACL_CHECK(aclrtFree(d_BPointerPadding));

    ACL_CHECK(aclrtFree(d_AicAivWorkspace));
    ACL_CHECK(aclrtFree(d_AicAivWorkspacePointer));

    INFO_LOG("Kernel host was executed successfully. ");

}